import { nanoid } from 'nanoid/non-secure';
import { StructuredNode, SendRequest, Turn, TurnStatus, NodeType, ToolUseNode, AgentMemoryNode } from './types';
import { ModelProvider } from './transport';
import { ToolRegistry } from './tools';
import { InMemoryMemoryStore, MemoryItem } from './memory';

export interface SessionEvents {
  onNode?: (turnId: string, node: StructuredNode) => void;
  onTurnStatusChange?: (turnId: string, status: TurnStatus, turn: Turn) => void;
  onError?: (turnId: string, error: Error) => void;
}

export interface SessionDeps {
  model: ModelProvider;
  tools?: ToolRegistry;
  memories?: InMemoryMemoryStore;
  allowAutoTools?: boolean;
  events?: SessionEvents;
}

export class Session {
  private turns = new Map<string, Turn>();

  constructor(private readonly cfg: SessionDeps) {}

  listTurns(): Turn[] { return [...this.turns.values()]; }

  getTurn(id: string): Turn | undefined { return this.turns.get(id); }

  async send(req: SendRequest): Promise<string> {
    const requestId = nanoid();
    const turnId = requestId;
    const turn: Turn = {
      id: turnId,
      requestId,
      status: TurnStatus.Sent,
      requestMessage: req.requestMessage,
      nodes: [],
    };
    this.turns.set(turnId, turn);

    (async () => {
      try {
        for await (const node of this.cfg.model.generate({ ...req, requestId })) {
          this.pushNode(turnId, node);

          // Handle tool execution
          if (node.type === NodeType.ToolUse && this.cfg.tools && this.cfg.allowAutoTools) {
            await this.executeToolNode(turnId, node as ToolUseNode);
          }

          // Handle memory nodes
          if (node.type === NodeType.AgentMemory && this.cfg.memories) {
            this.handleMemoryNode(node as AgentMemoryNode);
          }

          if (node.type === NodeType.MainTextFinished) {
            this.finishTurn(turnId, TurnStatus.Success);
          }
        }
      } catch (err) {
        const error = err instanceof Error ? err : new Error(String(err));
        this.cfg.events?.onError?.(turnId, error);
        this.finishTurn(turnId, TurnStatus.Failed);
      }
    })();

    return turnId;
  }

  cancel(turnId: string) {
    const turn = this.turns.get(turnId);
    if (turn && turn.status === TurnStatus.Sent) {
      this.finishTurn(turnId, TurnStatus.Cancelled);
      this.cfg.model.cancel?.(turn.requestId);
    }
  }

  async resendTurn(turnId: string): Promise<string> {
    const originalTurn = this.turns.get(turnId);
    if (!originalTurn) {
      throw new Error(`Turn ${turnId} not found`);
    }

    // Create new turn with same request message
    return this.send({
      requestMessage: originalTurn.requestMessage,
      // Could preserve other context from original request
    });
  }

  // Helper to get aggregated text from RAW_RESPONSE nodes
  getFullText(turnId: string): string {
    const turn = this.turns.get(turnId);
    if (!turn) return '';

    return turn.nodes
      .filter(n => n.type === NodeType.RawResponse)
      .map(n => (n as any).text)
      .join('');
  }

  // Helper to get suggested questions
  getSuggestedQuestions(turnId: string): string[] {
    const turn = this.turns.get(turnId);
    if (!turn) return [];

    const suggestedNode = turn.nodes.find(n => n.type === NodeType.SuggestedQuestions);
    return suggestedNode ? (suggestedNode as any).questions : [];
  }

  // Helper to get tool results
  getToolResults(turnId: string): Array<{ tool: string; input: any; result?: any; error?: string }> {
    const turn = this.turns.get(turnId);
    if (!turn) return [];

    return turn.nodes
      .filter(n => n.type === NodeType.ToolUse)
      .map(n => {
        const toolNode = n as any;
        return {
          tool: toolNode.tool,
          input: toolNode.input,
          result: toolNode.result,
          error: toolNode.error
        };
      });
  }

  private pushNode(turnId: string, node: StructuredNode) {
    const turn = this.turns.get(turnId);
    if (!turn) return;
    turn.nodes.push(node);

    // Emit node event
    this.cfg.events?.onNode?.(turnId, node);
  }

  private finishTurn(turnId: string, status: TurnStatus) {
    const turn = this.turns.get(turnId);
    if (!turn) return;
    const oldStatus = turn.status;
    turn.status = status;

    // Emit status change event if status actually changed
    if (oldStatus !== status) {
      this.cfg.events?.onTurnStatusChange?.(turnId, status, turn);
    }
  }

  private async executeToolNode(turnId: string, toolNode: ToolUseNode) {
    if (!this.cfg.tools) return;

    const tool = this.cfg.tools.get(toolNode.tool);
    if (!tool) {
      // Push error result if tool not found
      this.pushNode(turnId, {
        type: NodeType.ToolUse,
        ts: Date.now(),
        requestId: toolNode.requestId,
        tool: toolNode.tool,
        input: toolNode.input,
        error: `Tool '${toolNode.tool}' not found in registry`
      });
      return;
    }

    try {
      const result = await tool.run(toolNode.input, {
        log: (msg: string) => console.log(`[Tool ${toolNode.tool}]:`, msg)
      });

      // Push successful result
      this.pushNode(turnId, {
        type: NodeType.ToolUse,
        ts: Date.now(),
        requestId: toolNode.requestId,
        tool: toolNode.tool,
        input: toolNode.input,
        result
      });
    } catch (error) {
      // Push error result
      this.pushNode(turnId, {
        type: NodeType.ToolUse,
        ts: Date.now(),
        requestId: toolNode.requestId,
        tool: toolNode.tool,
        input: toolNode.input,
        error: error instanceof Error ? error.message : String(error)
      });
    }
  }

  private handleMemoryNode(memoryNode: AgentMemoryNode) {
    if (!this.cfg.memories) return;

    const memoryItem: MemoryItem = {
      id: memoryNode.memoryId,
      content: memoryNode.content,
      state: 'pending',
      requestId: memoryNode.requestId,
      timestamp: memoryNode.ts
    };

    this.cfg.memories.upsert(memoryItem);
  }
}

